/******************************************************
 * @file cnn_operation.c
 * @author Destiny 
 * @brief 卷积操作API定义
 * 
 * @version 0.1
 * @date 2024-07-18
 *****************************************************/
#include <malloc.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <time.h>
#include <math.h>
#include "cnn_operation.h"


/*本地函数声明*/;
static float KernelConvolution(PT_InoutputData ptInputData, PT_KernelParams ptKernelParams, int iCurRow, int iCurCol, int iCurNum,int iStrides);
static void Padding(PT_InoutputData ptInputData, PT_InoutputData ptPadData, int padding);
static float VectorMultiplication(PT_InoutputData ptInputData, PT_KernelParams ptKernelParams, int iCurNum);
static float MaxPoolOperation(PT_InoutputData ptInputData, int iCurRow, int iCurCol, int iCurDim, int iPoolSize, int iStrides);
static float AveragePoolOperation(PT_InoutputData ptInputData, int iCurRow, int iCurCol, int iCurDim, int iPoolSize, int iStrides);
/*o = [( iCol + iPadding -iKernelSize) / iStrides] + 1*/

/**
 * @brief 卷积核卷积操作
 * 
 * @param ptInputData 输入数据
 * @param ptKernelParams 卷积核参数
 * @param iCurRow 卷积起始位置 行
 * @param iCurCol 列
 * @param iCurNum 维度
 * @param iStrides 卷积步长
 * @return float 卷积结果
 */
static float KernelConvolution(PT_InoutputData ptInputData, PT_KernelParams ptKernelParams, int iCurRow, int iCurCol, int iCurNum,int iStrides)
{
    int i, j, k;
    float fSum =  0;
    for (k = 0; k < ptKernelParams->iDim; k++)
    {
        for (i = 0; i < ptKernelParams->iRow; i++)
        {
            for (j = 0; j < ptKernelParams->iCol; j++)
            {
                fSum += ptInputData->pppfDatas[k][i + iCurRow*iStrides][j + iCurCol*iStrides] * ptKernelParams->ppppfKernelParams[i][j][k][iCurNum];
            }
        }
    }
    return fSum + ptKernelParams->pfKernelBias[iCurNum];
}

/**
 * @brief 卷积层卷积操作
 * 
 * @param ptInputData 卷积层输入数据
 * @param ptKernelParams 卷积核参数
 * @param iStrides 步长
 * @param iPadding 填充数
 * @return PT_InoutputData 卷积结果（malloc分配内存）
 */
PT_InoutputData Convolution2D(PT_InoutputData ptInputData, PT_KernelParams ptKernelParams, int iStrides, int iPadding)
{
    int iCurNum, iCurRow, iCurCol;
    PT_InoutputData ptTemp;
    /*输出数据结构体构造*/
    int iOutputSize = (ptInputData->iCol + 2 * iPadding - ptKernelParams->iRow) / iStrides + 1;
    PT_InoutputData ptOutputData = (PT_InoutputData)malloc(sizeof(T_InoutputData));
    ptOutputData->iCol = iOutputSize;
    ptOutputData->iRow = iOutputSize;
    ptOutputData->iDim = ptKernelParams->iNum;
    MallocInoutputDatas(ptOutputData);
    if (iPadding)
    {
        PT_InoutputData ptPadData = malloc(sizeof(T_InoutputData));
        while (ptPadData == NULL)
        {
            DBG_PRINTF(DLOG_LVL_ERROR,DLOG_TAG,"ptPadData malloc faild!\r\n");
            ptPadData = malloc(sizeof(T_InoutputData));
        }
        Padding(ptInputData, ptPadData, iPadding); /*填充*/
        ptTemp = ptPadData;
    }
    else
    {
        ptTemp = ptInputData;
    }
    for (iCurNum = 0; iCurNum < ptOutputData->iDim; iCurNum++) /*逐个卷积核*/
    {
        for (iCurRow = 0; iCurRow < ptOutputData->iRow; iCurRow += iStrides) /*行*/
        {
            for (iCurCol = 0; iCurCol < ptOutputData->iCol; iCurCol += iStrides) /*列*/
            {
                ptOutputData->pppfDatas[iCurNum][iCurRow][iCurCol] = KernelConvolution(ptTemp, ptKernelParams,
                                                                                       iCurRow, iCurCol, iCurNum,iStrides);
            }
        }
    }
    FreeInoutputDatas(ptTemp);
    ptTemp = NULL;
    return ptOutputData;
}

/**
 * @brief Flatten操作
 * 
 * @param ptInputData 输入原始数据
 * @return PT_InoutputData 扁平化后数据（malloc分配）
 */
PT_InoutputData Flatten(PT_InoutputData ptInputData)
{
    int i, j, k;
    int iIndex ;
    PT_InoutputData ptOutputData = (PT_InoutputData)malloc(sizeof(T_InoutputData));
    ptOutputData->iDim = 1;
    ptOutputData->iRow = 1;
    ptOutputData->iCol = ptInputData->iCol * ptInputData->iRow * ptInputData->iDim; /*展开后的数据长度*/
    MallocInoutputDatas(ptOutputData);

    for (i = 0; i < ptInputData->iDim; i++)
    {
        for (j = 0; j < ptInputData->iRow; j++)
        {
            for (k = 0; k < ptInputData->iCol; k++)
            {
                iIndex = (i + j * (ptInputData->iCol*ptInputData->iDim) + k *ptInputData->iDim);    /*需转换输入参数iRow iCol iDim*/
                ptOutputData->pppfDatas[0][0][iIndex] = ptInputData->pppfDatas[i][j][k];
            }
        }
    }
    FreeInoutputDatas(ptInputData);
    return ptOutputData;
}

/**
 * @brief 向量乘法
 * 
 * @param ptInputData 输入数据
 * @param ptKernelParams 全连接层参数
 * @param iCurNum 当前维数
 * @return float 结果
 */
static float VectorMultiplication(PT_InoutputData ptInputData, PT_KernelParams ptKernelParams, int iCurNum)
{
    int iIndex;
    float fSum = 0;
    for (iIndex = 0; iIndex < ptInputData->iCol; iIndex++)
    {
        fSum += ptInputData->pppfDatas[0][0][iIndex] * ptKernelParams->ppppfKernelParams[iIndex][iCurNum][0][0];
    }
    return fSum + ptKernelParams->pfKernelBias[iCurNum];
}

/**
 * @brief 全连接层操作
 * 
 * @param ptInputData 输入数据
 * @param ptKernelParams 全连接层参数
 * @return PT_InoutputData 全链接输出
 */
PT_InoutputData Dense(PT_InoutputData ptInputData,PT_KernelParams ptKernelParams)
{
    int iCurNum;
    PT_InoutputData ptOutputData = (PT_InoutputData)malloc(sizeof(T_InoutputData));
    ptOutputData->iDim = 1;
    ptOutputData->iRow = 1;
    ptOutputData->iCol = ptKernelParams->iCol; /*展开后的数据长度*/
    MallocInoutputDatas(ptOutputData);
    for (iCurNum = 0; iCurNum < ptOutputData->iCol; iCurNum++)
    {
        ptOutputData->pppfDatas[0][0][iCurNum] = VectorMultiplication(ptInputData, ptKernelParams, iCurNum);
    }
    FreeInoutputDatas(ptInputData);
    return ptOutputData;
}

/**
 * @brief 最大池化操作
 * 
 * @param ptInputData 输入
 * @param ptKernelParams 卷积核参数
 * @param iCurRow 卷积起始位置 行
 * @param iCurCol 列
 * @param iCurDim 维度
 * @param iPoolSize 池化大小
 * @param iStrides 卷积步长
 * @return float 池化结果
 */
static float MaxPoolOperation(PT_InoutputData ptInputData, int iCurRow, int iCurCol, int iCurDim, int iPoolSize, int iStrides)
{
    int i, j;
    float fMax = ptInputData->pppfDatas[iCurDim][iCurRow*iStrides][iCurCol*iStrides];
    for (i = 0; i < iPoolSize; i++)
    {
        for (j = 0; j < iPoolSize; j++)
        {
            fMax = (fMax > ptInputData->pppfDatas[iCurDim][iCurRow * iStrides + i][iCurCol * iStrides + j]) ? fMax :\
                            ptInputData->pppfDatas[iCurDim][iCurRow * iStrides + i][iCurCol * iStrides + j];
        }
    }
    return fMax;
}

/**
 * @brief 平均池化
 * 
 * @param ptInputData 输入
 * @param ptKernelParams 卷积核参数
 * @param iCurRow 卷积起始位置 行
 * @param iCurCol 列
 * @param iCurDim 维度
 * @param iPoolSize 池化大小
 * @param iStrides 卷积步长
 * @return float 池化结果
 */
static float AveragePoolOperation(PT_InoutputData ptInputData, int iCurRow, int iCurCol, int iCurDim, int iPoolSize, int iStrides)
{
    int i, j;
    float fSum = 0;
    for (i = 0; i < iPoolSize; i++)
    {
        for (j = 0; j < iPoolSize; j++)
        {
            fSum += ptInputData->pppfDatas[iCurDim][iCurRow * iStrides + i][iCurCol * iStrides + j];
        }
    }
    return fSum/(iPoolSize*iPoolSize);
}

/**
 * @brief 池化操作，iPadding为0
 * 
 * @param ptInputData 输入
 * @param iStrides 步长
 * @param iPadding 填充数
 * @param iPoolSize 池化大小
 * @param iPoolSelect 0：平均池化;1:最大池化
 * @return PT_InoutputData 池化结果
 */
PT_InoutputData PoolOperation(PT_InoutputData ptInputData, int iStrides, int iPadding, int iPoolSize, int iPoolSelect)
{
    float (*pFuncTemp)(PT_InoutputData ptInputData, int iCurRow, int iCurCol, int iCurDim, int iPoolSize, int iStrides);
    if(iPoolSelect)
    {
        pFuncTemp = MaxPoolOperation;
    }
    else
    {
        pFuncTemp = AveragePoolOperation;
    }
    int iCurDim, iCurRow, iCurCol;
    PT_InoutputData ptTemp;
    /*输出数据结构体构造*/
    int iOutputSize = (ptInputData->iCol + 2 * iPadding - iPoolSize) / iStrides + 1;
    PT_InoutputData ptOutputData = (PT_InoutputData)malloc(sizeof(T_InoutputData));
    ptOutputData->iCol = iOutputSize;
    ptOutputData->iRow = iOutputSize;
    ptOutputData->iDim = ptInputData->iDim;
    MallocInoutputDatas(ptOutputData);
    if (iPadding)
    {
        PT_InoutputData ptPadData = malloc(sizeof(T_InoutputData));
        while (ptPadData == NULL)
        {
            DBG_PRINTF(DLOG_LVL_ERROR,DLOG_TAG,"ptPadData malloc faild!\r\n");
            ptPadData = malloc(sizeof(T_InoutputData));
        }
        Padding(ptInputData, ptPadData, iPadding); /*填充*/
        ptTemp = ptPadData;
    }
    else
    {
        ptTemp = ptInputData;
    }
    for (iCurDim = 0; iCurDim < ptOutputData->iDim; iCurDim++) /*逐通道*/
    {
        for (iCurRow = 0; iCurRow < ptOutputData->iRow; iCurRow++) /*行*/
        {
            for (iCurCol = 0; iCurCol < ptOutputData->iCol; iCurCol++) /*列*/
            {
                ptOutputData->pppfDatas[iCurDim][iCurRow][iCurCol] = pFuncTemp(ptTemp,iCurRow,iCurCol,iCurDim,iPoolSize,iStrides);
            }
        }
    }
    FreeInoutputDatas(ptTemp);
    ptTemp = NULL;
    return ptOutputData;
}

/**
 * @brief relu激活函数
 * 
 * @param ptInputData 输入
 */
void ActivationRelu(PT_InoutputData ptInputData)
{
    int iCurDim, iCurRow, iCurCol;
    for (iCurDim = 0; iCurDim < ptInputData->iDim; iCurDim++) /*逐通道*/
    {
        for (iCurRow = 0; iCurRow < ptInputData->iRow; iCurRow++) /*行*/
        {
            for (iCurCol = 0; iCurCol < ptInputData->iCol; iCurCol++) /*列*/
            {
                if(ptInputData->pppfDatas[iCurDim][iCurRow][iCurCol] > 0)
                {
                    continue;
                }
                ptInputData->pppfDatas[iCurDim][iCurRow][iCurCol] = 0;
            }
        }
    }
}

/**
 * @brief Tanh激活函数
 * 
 * @param ptInputData 输入
 */
void ActivationTanh(PT_InoutputData ptInputData)
{
    float fTemp;
    int iCurDim, iCurRow, iCurCol;
    for (iCurDim = 0; iCurDim < ptInputData->iDim; iCurDim++) /*逐通道*/
    {
        for (iCurRow = 0; iCurRow < ptInputData->iRow; iCurRow++) /*行*/
        {
            for (iCurCol = 0; iCurCol < ptInputData->iCol; iCurCol++) /*列*/
            {
                fTemp = ptInputData->pppfDatas[iCurDim][iCurRow][iCurCol];
                ptInputData->pppfDatas[iCurDim][iCurRow][iCurCol] = (exp(2 * fTemp) - 1) / (exp(2 * fTemp) + 1);
            }
        }
    }
}

/**
 * @brief softmax激活函数
 * 
 * @param ptInputData 输入
 */
void ActivationSoftmax(PT_InoutputData ptInputData) 
{
    float fMax = 0;
    float fSum = 0.0;
    int i;
    for (i = 0; i < ptInputData->iCol; i++)
    {
        fMax = (fMax > ptInputData->pppfDatas[0][0][i]) ? fMax : ptInputData->pppfDatas[0][0][i];
    }

    for (i = 0; i < ptInputData->iCol; i++)
    {
        ptInputData->pppfDatas[0][0][i] = exp(ptInputData->pppfDatas[0][0][i] - fMax); // 减去最大值，防止指数溢出
        fSum += ptInputData->pppfDatas[0][0][i];
    }
    for (i = 0; i < ptInputData->iCol; i++)
    {
         ptInputData->pppfDatas[0][0][i] /= fSum;
    }
}

/**
 * @brief Padding操作
 * 
 * @param ptInputData 原始输入数据
 * @param ptPadData 填充后的数据
 * @param padding 填充数
 */
static void Padding(PT_InoutputData ptInputData, PT_InoutputData ptPadData, int padding)
{
    int i, j, k;
    ptPadData->iDim = ptInputData->iDim;
    ptPadData->iRow = ptInputData->iRow + 2 * padding;
    ptPadData->iCol = ptInputData->iCol + 2 * padding;
    MallocInoutputDatas(ptPadData);
    for (i = 0; i < ptPadData->iDim; i++)
    {
        for (j = 0; j < ptPadData->iRow; j++)
        {
            for (k = 0; k < ptPadData->iCol; k++)
            {
                if (k < padding || k >= (ptInputData->iCol + padding) || j < padding || j >= (ptInputData->iRow + padding))
                {
                    ptPadData->pppfDatas[i][j][k] = 0;
                }
                else
                {
                    ptPadData->pppfDatas[i][j][k] = ptInputData->pppfDatas[i][j - padding][k - padding];
                }
            }
        }
    }
    FreeInoutputDatas(ptInputData);
}

/**
 * @brief dropout 函数
 * 
 * @param ptInputData 输入
 * @param fDropoutRate dropout比率
 */
void Dropout(PT_InoutputData ptInputData, float fDropoutRate) 
{
    int i;
    float fRandomValue;
    srand(time(NULL));  // 设置随机数生成器的种子
    for (i = 0; i < ptInputData->iCol; i++) 
    {
        fRandomValue = (float)rand() / RAND_MAX;  // 生成 0 到 1 之间的随机值
        if (fRandomValue < fDropoutRate) 
        {
            ptInputData->pppfDatas[0][0][i] = 0;  // 将某些元素设置为 0
        }
    }
}

/**
 * @brief 为ptInoutputData分配内存
 * 
 * @param ptInoutputData 输入数据结构体
 */
void MallocInoutputDatas(PT_InoutputData ptInoutputData)
{
    int i, j;
    ptInoutputData->pppfDatas = (float ***)malloc(sizeof(float **) * ptInoutputData->iDim);
    while (ptInoutputData->pppfDatas == NULL)
    {
        DBG_PRINTF(DLOG_LVL_ERROR,DLOG_TAG,"pppfDatas malloc  faild!\r\n");
        ptInoutputData->pppfDatas = (float ***)malloc(sizeof(float **) * ptInoutputData->iDim);
    }
    for (i = 0; i < ptInoutputData->iDim; i++)
    {
        ptInoutputData->pppfDatas[i] = (float **)malloc(sizeof(float *) * ptInoutputData->iRow);
        while (ptInoutputData->pppfDatas[i] == NULL)
        {
            DBG_PRINTF(DLOG_LVL_ERROR,DLOG_TAG,"pppfDatas[%d] malloc  faild!\r\n", i);
            ptInoutputData->pppfDatas[i] = (float **)malloc(sizeof(float *) * ptInoutputData->iRow);
        }
        for (j = 0; j < ptInoutputData->iRow; j++)
        {
            ptInoutputData->pppfDatas[i][j] = (float *)malloc(sizeof(float) * ptInoutputData->iCol);
            while (ptInoutputData->pppfDatas[i][j] == NULL)
            {
                DBG_PRINTF(DLOG_LVL_ERROR,DLOG_TAG,"pppfDatas[%d][%d] malloc  faild!\r\n", i, j);
                ptInoutputData->pppfDatas[i][j] = (float *)malloc(sizeof(float) * ptInoutputData->iCol);
            }
        }
    }
}

/**
 * @brief 释放为ptInoutputData分配的内存
 * 
 * @param ptInoutputData 输入数据结构体
 */
void FreeInoutputDatas(PT_InoutputData ptInoutputData)
{
    int i,j;
    for (i = 0; i < ptInoutputData->iDim; i++)
    {
        for (j = 0; j < ptInoutputData->iRow; j++)
        {
            free(ptInoutputData->pppfDatas[i][j]);
        }
        free(ptInoutputData->pppfDatas[i]);
    }
    free(ptInoutputData->pppfDatas);
    free(ptInoutputData);
    ptInoutputData = NULL;
}